import gradio as gr
from PIL import Image, ImageOps, ImageChops
from utils import generate_random_dir, generate_random_file
from utils.ps_canvas import PSTools
from body_seg import BodySeg, body_part
import os

body_model = BodySeg()


def main(origin_image: Image, body: list) -> Image:
    # 设置数据缓存目录
    _save_dir = generate_random_dir()
    # 数据先保存
    origin_filename = generate_random_file(dir_=_save_dir)
    origin_image.save(origin_filename)
    # 开始人体分割检测
    body_model.seg_target=body
    mask_filename = body_model.get_mask(
        image_path=origin_filename,
        save_dir=_save_dir)
    PSTools.find_max_mask(
        mask_filename, out_path=mask_filename, bitwise_not=True)  # 找最大的脸
    PSTools.mask_remove_isolated_pixels(
        mask_filename, out_path=mask_filename)  # 修复孤立点
    filename = PSTools.compare_mask_and_origin(
        origin_filename, mask_filename, _save_dir)  # mask与原像素差值

    image = Image.open(filename)
    image1 = Image.open(mask_filename)
    os.system(f"rm -rf {_save_dir}")

    return image1, image


def demo():
    with gr.Row():
        with gr.Column(scale=1):
            origin_image = gr.Image(type="pil", label="原始图像")
            body = gr.Dropdown(choices=body_part,
                               multiselect=True, label="身体部件")
        out_image = gr.Image(type="pil", label="分割目标图")
        mask_out_image = gr.Image(type="pil", label="分割蒙版图")
    gr.Button("运行").click(main, inputs=[
        origin_image, body], outputs=[out_image, mask_out_image])
